package com.huawei.aiservice.utils;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.huawei.aiservice.entity.dto.ChatRequest;
import com.huawei.aiservice.entity.dto.ChatResponse;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;

import java.util.Arrays;
import java.util.concurrent.TimeUnit;

@Slf4j
public class AIUtils {

    private static final String modelName = "qwen2-72b-instruct";

    private static final String url = "http://1.95.90.174:8080/v1/chat/completions";

    public static String getResponse(String systemContent, String userContent) throws Exception {
        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(5, TimeUnit.MINUTES)  // 设置连接超时为30秒
                .writeTimeout(5, TimeUnit.MINUTES)    // 设置写入超时为30秒
                .readTimeout(5, TimeUnit.MINUTES)     // 设置读取超时为30秒
                .build();
        // 请求体
        MediaType mediaType = MediaType.parse("application/json");
        ChatRequest chatRequest = getChatRequest(systemContent, userContent);
        ObjectMapper objectMapper = new ObjectMapper();
        String jsonString = objectMapper.writeValueAsString(chatRequest);
        RequestBody body = RequestBody.create(jsonString, mediaType);

        // 创建请求
        Request request = new Request.Builder()
                .url(url)
                .post(body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
//                .addHeader("Authorization", "Bearer "+token)
                .build();

        Response response = client.newCall(request).execute();
        if (response.isSuccessful() && response.code()==200) {
            ChatResponse chatResponse = objectMapper.readValue(response.body().string(), ChatResponse.class);
            return chatResponse.getChoices().get(0).getMessage().getContent();
        } else {
            log.error("Request failed with response code: " + response.code()+" "+response.body().string());
            return null;
        }
    }

    private static ChatRequest getChatRequest(String systemContent, String userContent) {
        ChatRequest chatRequest = new ChatRequest();
        chatRequest.setModel(modelName);
        chatRequest.setStream(false);
        ChatRequest.Message systemRole = new ChatRequest.Message();
        systemRole.setRole("system");
        systemRole.setContent(systemContent);
        ChatRequest.Message userRole = new ChatRequest.Message();
        userRole.setRole("user");
        userRole.setContent(userContent);
        chatRequest.setMessages(Arrays.asList(systemRole, userRole));
        return chatRequest;
    }
}
